package org.example.test.test;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;


/**
 * 给定一个非空二叉树的根节点 root , 以数组的形式返回每一层节点的平均值。
 */
public class AverageOfLevels {
    HashMap<Integer, List<Integer>> map = new HashMap<>();

    public static void main(String[] args) {

    }

    public List<Double> averageOfLevels(TreeNode root) {
        dfs(root, 0);
        ArrayList<Double> doubles = new ArrayList<>();
        for (int i = 0; i < map.size(); i++) {
            List<Integer> list = map.get(i);
            double sum = list.stream().mapToDouble(Integer::doubleValue).sum();
            doubles.add(sum / list.size());
        }
        return doubles;
    }

    /**
     * 将二叉树的节点按照层级存放在Map中
     *
     * @param root  当前节点
     * @param index 当前节点所在的层级
     */
    public void dfs(TreeNode root, int index) {
        if (root == null) {
            return;
        }
        List<Integer> indexList = map.getOrDefault(index, new ArrayList<>());
        indexList.add(root.val);
        map.put(index, indexList);
        dfs(root.left, index + 1);
        dfs(root.right, index + 1);
    }
}
